import pandas as pd
import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import matplotlib.gridspec as gridspec
from matplotlib.figure import Figure

import seaborn as sns # 热力图
import itertools
import datetime
import matplotlib.pyplot as plt
import statsmodels
import statsmodels.api as sm
from statsmodels.tsa.stattools import adfuller # ADF检验
from statsmodels.stats.diagnostic import acorr_ljungbox # 白噪声检验
from statsmodels.graphics import utils
from statsmodels.tsa.stattools import acf, pacf
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf # 画图定阶
from statsmodels.tsa.arima_model import ARIMA # ARIMA模型

class Arima_model(FigureCanvas):
    def __init__(self, width, heigh, dpi=100):
        self.figs = Figure(figsize=(width, heigh), dpi=dpi) #创建一个Figure
        super(Arima_model, self).__init__(self.figs)
        self.axes = self.figs.add_subplot(111)

    def data_resampling(self, file, key, part=0.75):
        df = pd.read_csv(file, parse_dates=True, index_col = "Date")
        len_ = int(part*len(df) + 1)
        train_data = df.iloc[:len_, :]
        test_data = df.iloc[len_+1:, :]
        
        # 画图
        self.axes.set_title('data')
        self.axes.plot(train_data[key], color='pink')

        return train_data, test_data

    def data_stationarity(self, data, d=1):
        # 差分法 对时间序列进行平滑处理
        for i in range(d):
            diff_data = data.diff(1).fillna(0.0)
        self.axes.set_title('diff')
        self.axes.plot(diff_data, color='blue')
        return diff_data

    def ADF_test(self, data):
        # 平稳性检验 ADF
        x = np.array(data)
        result = adfuller(x, autolag=='AIC')

        # 对比adf的结果是否小于10%时的假设检验 以及P-value是否非常接近0（P-value越小越好）
        if result[0] < result[4]["1%"] and result[1] < 100**(-8):
            return True
        else:
            return False

    def random_test(self, data):
        # 随机性检验 （白噪声检验）
        p_value = acorr_ljungbox(data, lags=1) # 二维数组，第二维是p值
        # 显著性水平小于0.05，以95%的置信水平拒绝原假设，数据非白噪声序列
        if p_value[1] < 0.05:
            return True
        else:
            # 随机序列
            return False

    '''
    时间序列定阶 (p, q) 和模型选择
    1) ACF 和 PACF 利用截尾方法来确定（p, q）
    2) 信息准则定阶 (AIC、BIC) 选择模型 
    '''  
    def plt_acf(self, data, lags=40):
        plot_acf(data, lags=lags, ax=self.axes) # lags延迟数
        #plt.savefig('acf.jpg')

    def plt_pacf(self, data, lags=40):
        plot_pacf(data, lags=lags, ax=self.axes)# lags延迟数

    def determinante_AIC(self, data, ar=5, ma=4):
        data.index = pd.DatetimeIndex(data.index.values, freq=data.index.inferred_freq)
        AIC = sm.tsa.arma_order_select_ic(data, max_ar=ar, max_ma=ma, ic='aic')['aic_min_order']
        return AIC

    def determinante_BIC(self, data, ar=5, ma=4):
        data.index = pd.DatetimeIndex(data.index.values, freq=data.index.inferred_freq)
        BIC = sm.tsa.arma_order_select_ic(data, max_ar=ar, max_ma=ma, ic='bic')['bic_min_order']
        return BIC

    '''
    画出AIC、BIC的热力图
    '''
    def heatmap_AIC(self, data, p_min=0, p_max=5, d_min=1, d_max=3, q_min=0, q_max=5):

        results_aic = pd.DataFrame(index=['AR{}'.format(i) for i in range(p_min,p_max+1)],
                                columns=['MA{}'.format(i) for i in range(q_min,q_max+1)])
    
        for p,d,q in itertools.product(range(p_min,p_max+1),
                                    range(d_min,d_max+1),
                                    range(q_min,q_max+1)):
            if p==0 and d==0 and q==0:
                results_aic.loc['AR{}'.format(p), 'MA{}'.format(q)] = np.nan
                continue
            try:
                model = sm.tsa.ARIMA(data, order=(p, d, q),
                                #enforce_stationarity=False,
                                #enforce_invertibility=False,
                                )
                results = model.fit()
                results_aic.loc['AR{}'.format(p), 'MA{}'.format(q)] = results.aic
            except:
                continue
        results_aic = results_aic[results_aic.columns].astype(float) 

        sns.heatmap(results_aic,
                    # mask=results_bic.isnull(),
                    ax=self.axes,
                    annot=True,
                    fmt='.2f',
                    )
        self.axes.set_title('AIC')

    def heatmap_BIC(self, data, p_min=0, p_max=5, d_min=1, d_max=3, q_min=0, q_max=5):
    
        results_bic = pd.DataFrame(index=['AR{}'.format(i) for i in range(p_min,p_max+1)],
                                columns=['MA{}'.format(i) for i in range(q_min,q_max+1)])
    
        for p,d,q in itertools.product(range(p_min,p_max+1),
                                    range(d_min,d_max+1),
                                    range(q_min,q_max+1)):
            if p==0 and d==0 and q==0:
                results_bic.loc['AR{}'.format(p), 'MA{}'.format(q)] = np.nan
                continue
            try:
                model = sm.tsa.ARIMA(data, order=(p, d, q),
                                #enforce_stationarity=False,
                                #enforce_invertibility=False,
                                )
                results = model.fit()
                results_bic.loc['AR{}'.format(p), 'MA{}'.format(q)] = results.bic
            except:
                continue
        results_bic = results_bic[results_bic.columns].astype(float) 

        sns.heatmap(results_bic,
                    # mask=results_bic.isnull(),
                    ax=self.axes,
                    annot=True,
                    fmt='.2f',
                    )
        self.axes.set_title('BIC')

    def in_sample_prediction(self, train_data, order):
        # 样本内预测
        arima_model = sm.tsa.ARIMA(endog=train_data, order=order) # ARIMA模型
        arima = arima_model.fit(disp=0) # 激活模型

        # arima.plot_predict()
        # plt.legend(labels=['forecast', 'actual'], loc='upper left', fontsize=10)
        # plt.title('In Sample Prediction')
        # plt.show()

        arima.plot_predict(ax=self.axes)
        self.axes.set_title('In Sample Prediction')
        self.axes.legend(labels=['forecast', 'actual'], loc='upper left', fontsize=7)

    def out_sample_prediction(self, train_data, test_data, order):
        # 样本外预测（需要从 train_data 的一个时间开始，另一个时间结束）

        arima_model = sm.tsa.ARIMA(endog=train_data, order=order) # ARIMA模型
        arima = arima_model.fit(disp=0) # 激活模型

        fc, se, conf = arima.forecast(len(test_data), alpha=0.05)

        fc_data = pd.Series(fc, index=test_data.index)
        lower_data = pd.Series(conf[:, 0], index=test_data.index)
        upper_data = pd.Series(conf[:, 1], index=test_data.index)

        # plt.plot(train_data, color='#149514', label='training')
        # plt.plot(test_data, color='#1F77B4', label='actual')
        # plt.plot(fc_data, color='#FF7F0E', label='forecast')
        # plt.fill_between(lower_data.index, lower_data, upper_data, color='k', alpha=.15)
        # plt.legend(loc='upper left', fontsize=10)
        # plt.title('Out Sample Prediction')
        # plt.show()

        self.axes.plot(train_data, color='#149514', label='training')
        self.axes.plot(test_data, color='#1F77B4', label='actual')
        self.axes.plot(fc_data, color='#FF7F0E', label='forecast')
        self.axes.fill_between(lower_data.index, lower_data, upper_data, color='k', alpha=.15)
        self.axes.set_title('Out Sample Prediction')
        self.axes.legend(loc='upper left', fontsize=7)

        return fc_data
        
    '''
    模型评价:
    '''
    def model_evaluate(self, train_data, order):
        arima_model = sm.tsa.SARIMAX(endog=train_data, order=order) # SARIMAX模型
        arima = arima_model.fit() # 激活模型

        self.figs.add_subplot(221)
        arima.plot_diagnostics(fig=self.figs)
        
# 测试用代码
if __name__ == '__main__':
    file = 'data\\sentiment.csv'
    Primarykey = 'UMCSENT'
    ARIMA = Arima_model(100, 100)
    train_data, test_data = ARIMA.data_resampling(file, Primarykey, part=0.9) 
    diff_data = ARIMA.data_stationarity(train_data[[Primarykey]], d=1) 
    ARIMA.plt_acf(diff_data[[Primarykey]]) 
    ARIMA.plt_pacf(diff_data[[Primarykey]]) 
    AIC = ARIMA.determinante_AIC(diff_data[Primarykey])
    print(AIC) 
    BIC = ARIMA.determinante_BIC(diff_data[Primarykey])
    print(BIC)

    p, q = AIC
    p, d, q = 3, 1, 2
    order_ = (p, d, q)
        
    ARIMA.in_sample_prediction(train_data[[Primarykey]], order_)
    ARIMA.out_sample_prediction(train_data[[Primarykey]], test_data[[Primarykey]], order_)

    ARIMA.model_evaluate(train_data[[Primarykey]], order_) 
        